#include "ScatterVal.h"

#include <iostream>

#include "TBranch.h"
#include "TLeaf.h"
#include "TTreeReader.h"
#include "TTreeReaderArray.h"
#include "TTreeReaderValue.h"

std::unordered_map<std::string, int> ScatterVal::m_SupportType = {  {"Int_t", 0}, {"Float_t", 1}, {"Double_t", 2}, 
                                            {"int", 0}, {"float", 1}, {"double", 2}, 
                                            {"vector<Int_t>", 3}, {"vector<Float_t>", 4}, {"vector<Double_t>", 5}, 
                                            {"vector<int>", 3}, {"vector<float>", 4}, {"vector<double>", 5} };

ScatterVal::ScatterVal(std::string name, TTree* tree)
{
    TTreeReaderValue<int>* IntVal = nullptr;
    TTreeReaderValue<float>* FloatVal = nullptr;
    TTreeReaderValue<double>* DoubleVal = nullptr;

    TTreeReaderArray<int>* IntVec = nullptr;
    TTreeReaderArray<float>* FloatVec = nullptr;
    TTreeReaderArray<double>* DoubleVec = nullptr;

    m_OutVector.clear();

    std::string type = tree->GetBranch(name.c_str())->FindLeaf(name.c_str())->GetTypeName();
    TTreeReader treeReader(tree);

    m_Type = m_SupportType[type];
    switch(m_Type)
    {
        case 0:
        {
            IntVal = new TTreeReaderValue<int>(treeReader, name.c_str());
            break;
        }
        case 1:
        {
            FloatVal = new TTreeReaderValue<float>(treeReader, name.c_str());
            break;
        }
        case 2:
        {
            DoubleVal = new TTreeReaderValue<double>(treeReader, name.c_str());
            break;
        }
        case 3:
        {
            IntVec = new TTreeReaderArray<int>(treeReader, name.c_str());
            break;
        }
        case 4:
        {
            FloatVec = new TTreeReaderArray<float>(treeReader, name.c_str());
            break;
        }
        case 5:
        {
            DoubleVec = new TTreeReaderArray<double>(treeReader, name.c_str());
            break;
        }
        default:
        {
            std::cout << "Error: Data type: \" " << type << " \" NOT SUPPORT for Scatter Plot." << std::endl;
            std::cout << "Avalable type: ";
            for (auto type : m_SupportType)
            {
                std::cout << type.first << "\t";
            }
            exit(EXIT_FAILURE);
        }
        }

        Long64_t numOfEntries = tree->GetEntries();
        m_OutVector.resize(numOfEntries, 0.);
        for (Long64_t iE = 0; iE < numOfEntries; ++iE)
        {
            treeReader.Next();
            switch (m_Type)
            {
                case 0:
                {
                    m_OutVector[iE] = double(**IntVal);
                    break;
                }
                case 1:
                {
                    m_OutVector[iE] = double(**FloatVal);
                    break;
                }
                case 2:
                {
                    m_OutVector[iE] = double(**DoubleVal);
                    break;
                }
                case 3:
                {
                    m_OutVector[iE] = double((*IntVec)[0]);
                    break;
                }
                case 4:
                {
                    m_OutVector[iE] = double((*FloatVec)[0]);
                    break;
                }
                case 5:
                {
                    m_OutVector[iE] = double((*DoubleVec)[0]);
                    break;
                }
            }
        }
    switch (m_Type)
    {
        case 0:
        {
            delete IntVal;
            break;
        }
        case 1:
        {
            delete FloatVal;
            break;
        }
        case 2:
        {
            delete DoubleVal;
            break;
        }
        case 3:
        {
            delete IntVec;
            break;
        }
        case 4:
        {
            delete FloatVec;
            break;
        }
        case 5:
        {
            delete DoubleVec;
            break;
        }
    }
}

ScatterVal::~ScatterVal(){}